# coding=utf-8
# -*- coding: UTF-8 -*-
#
# This file is part of the kernelstudio package.
#
# (c) 2014-2025 zlin <admin@kernelstudio.com>
#
# For the full copyright and license information, please view the LICENSE file
# that was distributed with this source code.

import typing as t
from pathlib import Path

from ksai.document.context import Document, DocumentSection
from ksai.document.extractor.csv_extractor import CsvDocumentExtractor
from ksai.document.extractor.docx_extractor import DocxDocumentExtractor
from ksai.document.extractor.excel_extractor import ExcelDocumentExtractor
from ksai.document.extractor.extractor import DocumentExtractor
from ksai.document.extractor.html_extractor import HtmlDocumentExtractor
from ksai.document.extractor.image_ocr_extractor import ImageOcrDocumentExtractor
from ksai.document.extractor.json_extractor import JsonDocumentExtractor
from ksai.document.extractor.pdf_extractor import PdfDocumentExtractor
from ksai.document.extractor.plan_text_extractor import PlanTextDocumentExtractor
from ksai.document.extractor.pptx_extractor import PptxDocumentExtractor
from ksai.document.extractor.properties_extractor import PropertiesDocumentExtractor
from ksai.document.extractor.yaml_extractor import YamlDocumentExtractor
from ksai.persistence.storage.attachment import Attachment
from ksai.persistence.storage.metadata import read_metadata, write_metadata


class DocumentExtractorFactory:

    def __init__(self, extractors: t.List[DocumentExtractor]):
        super().__init__()
        self.extractors = extractors

    def get_extractor(self, attachment: Attachment = None) -> DocumentExtractor | None:
        if self.has_extractors():
            for extractor in self.extractors:
                if extractor.supports(attachment):
                    return extractor
        return None

    def extract(self, attachment: Attachment = None) -> Document:
        metadata = read_metadata(attachment)
        if metadata and metadata.get('__sections'):
            sections = [DocumentSection.from_dict(d=d) for d in metadata.get('__sections', [])]
            return Document(attachment, sections)

        extractor = self.get_extractor(attachment)
        sections = []
        if extractor is not None:
            sections = extractor.extract(attachment)
            metadata = read_metadata(attachment)
            if extractor.is_cacheable():
                metadata['__sections'] = sections
            write_metadata(attachment, metadata)
        return Document(attachment, sections)

    def extract_from_file(self, filepath: t.Union[str, Path] = None) -> Document:
        return self.extract(Attachment.from_file(filepath))

    def extract_content(self, attachment: Attachment = None) -> str:
        document = self.extract(attachment)
        if not document:
            return ''
        if document.has_sections():
            contents = [str(d.text) for d in document.sections]
            return '\n\n\n'.join(contents)
        return ''

    def extract_content_from_file(self, filepath: t.Union[str, Path] = None) -> str:
        return self.extract_content(Attachment.from_file(filepath))

    def has_extractors(self) -> bool:
        return len(self.extractors) > 0

    @property
    def supported_mines(self) -> t.List[str]:
        mimes = []
        if self.has_extractors():
            for extractor in self.extractors:
                supported_mimes = extractor.mimes()
                if supported_mimes:
                    mimes += supported_mimes
        return mimes

    @property
    def supported_extensions(self) -> t.List[str]:
        extensions = []
        if self.has_extractors():
            for extractor in self.extractors:
                supported_extensions = extractor.extensions()
                if supported_extensions:
                    extensions += supported_extensions
        return extensions


def create_extractor_factory() -> DocumentExtractorFactory:
    default_extractors = [
        JsonDocumentExtractor(),
        ImageOcrDocumentExtractor(),
        YamlDocumentExtractor(),
        PlanTextDocumentExtractor(),
        HtmlDocumentExtractor(),
        PropertiesDocumentExtractor(),
        DocxDocumentExtractor(),
        CsvDocumentExtractor(),
        ExcelDocumentExtractor(),
        PptxDocumentExtractor(),
        PdfDocumentExtractor(),
    ]
    return DocumentExtractorFactory(default_extractors)


extractor_factory = create_extractor_factory()
